from common.helperFunctions import distance, mean, std
import common.matfunc
import common.Simplex
import common.commonDB
from common.constants import *
from ptzmodels import *
import math

def convertLatLonToXYZ(corrs):
	"""
	Takes in a list of n correspondences, returns a list of x,y,z points in ENU space.
	Returns:
		x,y,z:   a list of n normalized coordinates.
		originLat,originLon,originAlt: the geographic location of the origin.
	"""
	originLat = corrs[0].latitude;
	originLon = corrs[0].longitude;
	originAlt = corrs[0].altitude;
	
	metersPerLat = distance(originLat, originLon, originLat+0.00001, originLon) / 0.00001;
	metersPerLon = distance(originLat, originLon, originLat, originLon+0.00001) / 0.00001;
	
	x = []
	y = []
	z = []
	for corr in corrs:
		x.append((corr.longitude - originLon) * metersPerLon);
		y.append((corr.latitude - originLat) * metersPerLat);
		z.append(corr.altitude - originAlt);
		
	return x,y,z,originLat,originLon,originAlt
	
def convertXYZToLatLonAlt(x,y,z,originLat,originLon,originAlt):
	"""
	Takes in a single (X,Y,Z) point in ENU space and the geographic location of the origin.
	Returns the geographic position of (X,Y,Z).
	"""
	
	latPerMeter =  0.00001 / distance(originLat, originLon, originLat+0.00001, originLon)
	lonPerMeter =  0.00001 / distance(originLat, originLon, originLat, originLon+0.00001)
	
	lat = latPerMeter * y + originLat
	lon = lonPerMeter * x + originLon
	alt = z + originAlt
	return lat, lon, alt

def makeProjectionMatrix(panRad, tiltRad, zoom, width, height, \
						tx, ty, tz, panOffset, tiltOffset, rollOffset, focalMin, focalMax):
	"""
	Input: intrinsic and extrinsic parameters of a camera given a single pan and tilt configuration.
	Output: The 3x4 projection matrix that satisfies these parameters.
	"""
	# K is the intrinsic camera matrix.
	# assuming a linear camera response...
	c = zoom
	focal = c * focalMax + (1-c) * focalMin
	aspectRatio = width / height;
	K = common.matfunc.Matrix([[focal * aspectRatio, 0, width / 2],
				[0, -focal, height / 2],
				[0, 0, 1]])
				
	# R is the rotation matrix.
	
	alpha = -(panRad + panOffset)
	beta  = tiltRad + tiltOffset
	gamma = rollOffset + math.pi
	
	ca = math.cos(alpha)
	cb = math.cos(beta)
	cg = math.cos(gamma)
	sa = math.sin(alpha)
	sb = math.sin(beta)
	sg = math.sin(gamma)
	R = common.matfunc.Matrix([[ca * cg - cb * sa * sg,  cg * sa + ca * cb * sg, sb * sg],
                        [-cb * cg * sa - ca * sg, ca * cb * cg - sa*sg,   cg * sb],
			            [sa * sb,                 -ca * sb,               cb]])
	# T is the translation matrix.
	center = common.matfunc.Matrix([[tx],[ty],[tz]])
	T = -R.mmul(center);
	
	# M is the full projection matrix.
	M = R.augment(T)
	M = K.mmul(M);
	
	return M
	

def makeManyProjectionMatrices(panNormals, tiltNormals, zoomNormals, width, height, \
                               tx, ty, tz, panOffset, tiltOffset, rollOffset, totalPanRad, totalTiltRad,\
							   focalMin, focalMax):
	"""
	Input: intrinsic and extrinsic parameters of a camera across many pan,tilt and zoom configurations.
	Output: the 3x4 projection matrices that satisfy these parameters.
	"""
	allMatrices = []
	n = len(panNormals)
	for i in range(n):
		panRad = panNormals[i] * totalPanRad;
		tiltRad = tiltNormals[i] * totalTiltRad;
		zoom = zoomNormals[i]
		mat = makeProjectionMatrix(panRad, tiltRad, zoom, width, height, tx, ty, tz, panOffset, tiltOffset, rollOffset, focalMin, focalMax)
		allMatrices.append(mat)
	return allMatrices

def ptzAngleError(x,y,z,u,v,panNormals,tiltNormals,zoomNormals,width, height, sol):
	"""
	Returns the sum of all angular errors.
	Input: x,y,z                              : positions in ENU space
	       u,v                                : positions in image coordinates
		   panNormals,tiltNormals,zoomNormals : pan, tilt, and zoom, in [0,1].
		   width, height                      : the width and height of the image.
		   sol                                : a 10-element solution vector, where
		      sol[0:2]=tx,ty,tz               :    the position of the camera in ENU space.
			  sol[3:5]=pan0,tilt0,roll0       :    the rotational offsets of the camera.
			  sol[6:7]=focalMin,focalMax      :    the minimum and maximum focal lengths of the camera.
			  sol[8:9]=dpan, dtilt            :    the total field of view of the camera in pan and tilt.
	Output:  the angular difference error.
	"""
	n = len(x)
	tx                = sol[0]
	ty                = sol[1]
	tz                = sol[2]
	panOffset         = sol[3]
	tiltOffset        = sol[4]
	#rollOffset        = sol[5]
	rollOffset        = 0
	focalMin          = sol[5]
	focalMax          = sol[6]
	totalPanRad       = sol[7]
	totalTiltRad      = sol[8]
	
	# assert that the focal lengths are both positive and that the minimum
	# is less than the maximum.
	if focalMin < 0 or focalMax < 0:
		return float('infinity')
	
	if focalMin > focalMax:
		return float('infinity')
		
	allM = makeManyProjectionMatrices(panNormals,tiltNormals,zoomNormals, \
	                                  width,height,   \
									  tx,ty,tz,   \
									  panOffset,tiltOffset, rollOffset, \
									  totalPanRad, totalTiltRad, \
									  focalMin,focalMax)

	w = -100
	camCenter = common.matfunc.Matrix([[tx],[ty],[tz]])
	error = 0
	for i in range(n):
		# from M and the image coordinates, find out which ray goes from the
		# camera center to the image through the user-provided image coordinates.
		
		# solving for XYZ:
		#   x*d = m11*X + m12*Y + m13*Z + m14
		#   y*d = m21*X + m22*Y + m23*Z + m24
		#     d = m31*X + m32*Y + m33*Z + m34
		#  (x*d - m14) = m11*X + m12*Y + m13*Z
		#  (y*d - m24) = m21*X + m22*Y + m23*Z 
		#     d - m34  = m31*X + m32*Y + m33*Z
		M = allM[i]
		lastColumn = common.matfunc.Matrix([[M[0][3]], 
		                     [M[1][3]],
							 [M[2][3]]])
		first3x3   = common.matfunc.Matrix([[M[0][0], M[0][1], M[0][2]], 
		                     [M[1][0], M[1][1], M[1][2]], 
							 [M[2][0], M[2][1], M[2][2]]])
		A = common.matfunc.Matrix([[w * u[i]],
		            [w * v[i]],
				    [w]])
		B = A - lastColumn
		
		projectedXYZ = first3x3.solve(B)
		projectedRay = projectedXYZ - camCenter
		mag = math.sqrt(projectedRay[0][0]**2 + projectedRay[1][0]**2 + projectedRay[2][0]**2)
		projectedRay = projectedRay / mag
		
		# now find the ray that goes from the camera center through the user-provided 3D XYZ coordinate.
		XYZPoint  = common.matfunc.Matrix([[x[i]],[y[i]],[z[i]]])
		rayXYZPoint = XYZPoint - camCenter
		mag = math.sqrt(rayXYZPoint[0][0]**2 + rayXYZPoint[1][0]**2 + rayXYZPoint[2][0]**2)
		rayXYZPoint = rayXYZPoint / mag
		
		# the error increases as the vectors become more and more opposite.
		#  i.e., error = 1 - ray1 (dot) ray2.
		error += 1 - projectedRay.tr().mmul(rayXYZPoint)[0][0]
		
	return error;

def calibratePTZScene(scene, iterations=PTZ_CALIBRATION_ITERATIONS):
	corrs = PTZCorrespondence.all().filter('scene =', scene).fetch(100)
	if len(corrs) < CORRESPONDENCES_REQUIRED_FOR_PTZ_CALIBRATION:
		if scene.calibrated:
			scene.cameraHeading         = None
			scene.cameraTilt            = None
			scene.cameraRoll            = None
			scene.cameraFocalLengthX    = None
			scene.cameraFocalLengthY    = None
			scene.calibrated = False
			scene.put()
		return
		
	panNormals, tiltNormals, zoomNormals = [], [], []
	imageUs, imageVs = [], []
	for corr in corrs:
		pan = corr.pan
		tilt = corr.tilt
		zoom = corr.zoom
		panNormal = (pan - scene.panLeftRange) / (scene.panRightRange - scene.panLeftRange)
		tiltNormal = (tilt - scene.tiltDownRange) / (scene.tiltUpRange - scene.tiltDownRange)
		zoomNormal = (zoom - scene.zoomOutRange) / (scene.zoomInRange - scene.zoomOutRange)
		panNormals.append(panNormal)
		tiltNormals.append(tiltNormal)
		zoomNormals.append(zoomNormal)
		imageUs.append(corr.u)
		imageVs.append(corr.v)
		
	x,y,z,originLat,originLon,originAlt	= convertLatLonToXYZ(corrs)
	
	initial = [mean(x), mean(y), mean(z), \
			  0, math.pi/2, \
			  scene.zoomInRange, scene.zoomOutRange, \
			  scene.totalPanDegrees * math.pi / 180, scene.totalTiltDegrees * math.pi / 180];
	stepSizes = [std(x)/4.0, std(y)/4.0, std(z)/4.0, \
			math.pi / 4, math.pi / 8, \
			50, 50, \
			math.pi / 16, math.pi / 16];
			
	calib = PTZCalibration.get_by_key_name('ptzcalib' + str(scene.index))
	if calib:
		initial[0] = calib.tx;
		initial[1] = calib.ty;
		initial[2] = calib.tz;
		initial[3] = calib.panOffsetRad;
		
		
	# sol is a 10-element vector, where
	#  sol[0:2] = tx,ty,tz, the position of the camera.
	#  sol[3:5] = p0, t0, r0, the pan, tilt, and roll offsets.
	#  sol[6:7] = focalMin, focalMax, the minimum and maximum focal lengths of the camera.
	#  sol[8:9] = panRange, tiltRange, the total field of view for pan and tilt, in radians.
	errFun = lambda sol: ptzAngleError(x,y,z,imageUs, imageVs,panNormals,tiltNormals,zoomNormals,scene.width, scene.height, sol)
			  
	s = common.Simplex.Simplex(errFun, initial, stepSizes)
	
	values, err, iter = s.minimize(epsilon= 0.00001 * len(corrs), maxiters=iterations, monitor=0)
	tx = values[0]
	ty = values[1]
	tz = values[2]
	panOffsetRad = values[3]
	tiltOffsetRad = values[4]
	#rollOffsetRad = values[5]
	rollOffsetRad = float(0.0)
	focalLengthMin = values[5]
	focalLengthMax = values[6]
	totalPanDegrees = values[7] * 180.0 / math.pi
	totalTiltDegrees = values[8] * 180.0 / math.pi
	
	cameraLat, cameraLon, cameraAlt = convertXYZToLatLonAlt(tx,ty,tz,originLat,originLon,originAlt)
	
	calib = common.commonDB.addOrModifyPTZCalibration(scene, originLat, originLon, originAlt, tx, ty, tz, panOffsetRad, \
											  tiltOffsetRad, rollOffsetRad, focalLengthMin, focalLengthMax);
	
	calib.averageError = err / float(len(corrs));
	calib.put()
	
	
	scene.cameraLatitude = cameraLat
	scene.cameraLongitude = cameraLon
	scene.cameraAltitude = cameraAlt
#	scene.totalPanDegrees = totalPanDegrees
#	scene.totalTiltDegrees = totalTiltDegrees
	scene.calibrated = True
	scene.put()